import math
import torch
from torch import nn
# from d2l import torch as d2l
# from tools import show_plt


debug = False


def sequence_mask(X, valid_len, value=0):
    """Mask irrelevant entries in sequences.

    Defined in :numref:`sec_seq2seq_decoder`"""
    maxlen = X.size(1)
    mask = torch.arange((maxlen), dtype=torch.float32,
                        device=X.device)[None, :] < valid_len[:, None]
    X[~mask] = value
    return X


def masked_softmax(X, valid_lens):
    """通过在最后一个轴上掩蔽元素来执行softmax操作"""
    # X:3D张量，valid_lens:1D或2D张量
    if valid_lens is None:
        return nn.functional.softmax(X, dim=-1)
    else:
        shape = X.shape
        if valid_lens.dim() == 1:
            _valid_lens = valid_lens
            valid_lens = torch.repeat_interleave(valid_lens, shape[1])
        else:
            valid_lens = valid_lens.reshape(-1)
        # 最后一轴上被掩蔽的元素使用一个非常大的负值替换，从而其softmax输出为0
        X = sequence_mask(X.reshape(-1, shape[-1]), valid_lens, value=-1e6)
        return nn.functional.softmax(X.reshape(shape), dim=-1)


class AdditiveAttention(nn.Module):
    """加性注意力"""
    def __init__(self, key_size, query_size, num_hiddens, dropout, **kwargs):
        super(AdditiveAttention, self).__init__(**kwargs)
        self.W_k = nn.Linear(key_size, num_hiddens, bias=False)
        self.W_q = nn.Linear(query_size, num_hiddens, bias=False)
        self.w_v = nn.Linear(num_hiddens, 1, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, queries, keys, values, valid_lens):
        q_0 = queries
        k_0 = keys
        v_0 = values
        if debug: print("shape 0 --- q_0, k_0, v_0:", q_0.shape, k_0.shape, v_0.shape)
        queries, keys = self.W_q(queries), self.W_k(keys)
        if debug: print("shape 1 --- q, k, v:", queries.shape, keys.shape, values.shape)
        # 在维度扩展后，
        # queries的形状：(batch_size，查询的个数，1，num_hidden)
        # key的形状：(batch_size，1，“键－值”对的个数，num_hiddens)
        # 使用广播方式进行求和
        if debug: print("shape 2 --- q_unsqueeze_2, k_unsqueeze_1:", queries.unsqueeze(2).shape, keys.unsqueeze(1).shape)
        features = queries.unsqueeze(2) + keys.unsqueeze(1)
        if debug: print("shape 2 --- features:", features.shape)

        features = torch.tanh(features)
        if debug: print("shape 3 --- tanh_features:", features.shape)

        # self.w_v仅有一个输出，因此从形 状中移除最后那个维度。
        # scores的形状：(batch_size，查询的个数，“键-值”对的个数)
        self.w_v(features).shape
        scores = self.w_v(features).squeeze(-1)
        if debug: print("shape 4 --- scores:", scores.shape)

        self.attention_weights = masked_softmax(scores, valid_lens)
        if debug: print("shape 5 --- self.attention_weights:", self.attention_weights.shape)

        # values的形状：(batch_size，“键－值”对的个数，值的维度)
        res = torch.bmm(self.dropout(self.attention_weights), values)
        if debug: print("shape 6 --- res:", res.shape)
        return res


class DotProductAttention(nn.Module):
    """缩放点积注意力"""
    def __init__(self, dropout, **kwargs):
        super(DotProductAttention, self).__init__(**kwargs)
        self.dropout = nn.Dropout(dropout)

    # queries的形状：(batch_size，查询的个数，d)
    # keys的形状：(batch_size，“键－值”对的个数，d)
    # values的形状：(batch_size，“键－值”对的个数，值的维度)
    # valid_lens的形状:(batch_size，)或者(batch_size，查询的个数)
    def forward(self, queries, keys, values, valid_lens=None):
        q_0 = queries
        k_0 = keys
        v_0 = values
        if debug: print("shape 0 --- q_0, k_0, v_0:", q_0.shape, k_0.shape, v_0.shape)

        old_shape = queries.shape

        flag__auto_reshape_4d = False
        if queries.dim() == 4 and keys.dim() == 4 and values.dim() == 4:
            flag__auto_reshape_4d = True

        if flag__auto_reshape_4d:
            queries = queries.reshape(old_shape[0], old_shape[1], old_shape[2] * old_shape[3])
            keys = keys.reshape(old_shape[0], old_shape[1], old_shape[2] * old_shape[3])
            values = values.reshape(old_shape[0], old_shape[1], old_shape[2] * old_shape[3])
            if debug: print("~~~ reshape 0 --- queries, keys, values:", queries.shape, keys.shape, values.shape)

        d = queries.shape[-1]
        # 设置transpose_b=True为了交换keys的最后两个维度
        # keys.transpose(1, 2).shape
        scores = torch.bmm(queries, keys.transpose(1, 2)) / math.sqrt(d)
        if debug: print("shape --- scores:", scores.shape)

        self.attention_weights = masked_softmax(scores, valid_lens)
        if debug: print("shape --- self.attention_weights:", self.attention_weights.shape)

        res = torch.bmm(self.dropout(self.attention_weights), values)
        if debug: print("shape --- res:", res.shape)

        if flag__auto_reshape_4d:
            res = res.reshape(old_shape)
        return res


def transpose_qkv(X, num_heads):
    """为了多注意力头的并行计算而变换形状"""
    # 输入X的形状: (batch_size，查询或者“键－值”对的个数，num_hiddens)

    # 输出X的形状: (batch_size，查询或者“键－值”对的个数，num_heads，num_hiddens/num_heads)
    num_hiddens = X.shape[-1]
    assert num_hiddens % num_heads == 0, "`num_hiddens`不能被`num_heads`整除?"
    X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)

    # 输出X的形状: (batch_size，num_heads，查询或者“键－值”对的个数, num_hiddens/num_heads)
    X = X.permute(0, 2, 1, 3)

    # 最终输出的形状: (batch_size*num_heads,查询或者“键－值”对的个数, num_hiddens/num_heads)
    return X.reshape(-1, X.shape[2], X.shape[3])


#@save
def transpose_output(X, num_heads):
    """逆转transpose_qkv函数的操作"""
    X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])
    X = X.permute(0, 2, 1, 3)
    return X.reshape(X.shape[0], X.shape[1], -1)


#@save
class MultiHeadAttention(nn.Module):
    """多头注意力"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 num_heads, dropout, bias=False, **kwargs):
        super(MultiHeadAttention, self).__init__(**kwargs)
        self.num_heads = num_heads
        self.attention = DotProductAttention(dropout)
        self.W_q = nn.Linear(query_size, num_hiddens, bias=bias)
        self.W_k = nn.Linear(key_size, num_hiddens, bias=bias)
        self.W_v = nn.Linear(value_size, num_hiddens, bias=bias)
        self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)

    def forward(self, queries, keys=None, values=None, valid_lens=None):
        # queries，keys，values的形状: (batch_size，查询或者“键－值”对的个数，num_hiddens)
        # valid_lens　的形状: (batch_size，)或(batch_size，查询的个数)
        # 经过变换后，输出的queries，keys，values　的形状:
        # (batch_size*num_heads，查询或者“键－值”对的个数， num_hiddens/num_heads)
        if keys is None:
            keys = queries
        if values is None:
            values = queries

        q_0, k_0, v_0 = queries, keys, values
        if debug: print('shape 0 --- [b, qkv, hiddens] --- q_0, k_0, v_0:', q_0.shape, k_0.shape, v_0.shape)

        old_shape = queries.shape

        flag__auto_reshape_4d = False
        if queries.dim() == 4 and keys.dim() == 4 and values.dim() == 4:
            flag__auto_reshape_4d = True

        if flag__auto_reshape_4d:
            queries = queries.reshape(old_shape[0], old_shape[1], old_shape[2] * old_shape[3])
            keys = keys.reshape(old_shape[0], old_shape[1], old_shape[2] * old_shape[3])
            values = values.reshape(old_shape[0], old_shape[1], old_shape[2] * old_shape[3])
            if debug: print("~~~ reshape 0 --- queries, keys, values:", queries.shape, keys.shape, values.shape)

        queries = transpose_qkv(self.W_q(queries), self.num_heads)
        keys = transpose_qkv(self.W_k(keys), self.num_heads)
        values = transpose_qkv(self.W_v(values), self.num_heads)
        if debug: print('shape 1 --- [b * heads, qkv, hiddens/heads] --- qkv:', queries.shape, keys.shape, values.shape)

        if valid_lens is not None:
            # 在轴0，将第一项（标量或者矢量）复制num_heads次，
            # 然后如此复制第二项，然后诸如此类。
            valid_lens = torch.repeat_interleave(valid_lens, repeats=self.num_heads, dim=0)

        # output的形状:(batch_size*num_heads，查询的个数，num_hiddens/num_heads)
        output = self.attention(queries, keys, values, valid_lens)

        if debug: print('shape 1 ------ output[b*heads, q, hiddens/heads]:', output.shape)

        # output_concat的形状:(batch_size，查询的个数，num_hiddens)
        output_concat = transpose_output(output, self.num_heads)
        if debug: print('shape 1 ------ output[b, q, hiddens]:', output.shape)

        res = self.W_o(output_concat)
        if flag__auto_reshape_4d:
            res = res.reshape(old_shape)
        if debug: print('shape 1 ------ W_o(output_concat)[b, q, hiddens]:', res.shape)
        return res


def get_total_parameters(model):
    # 遍历模型的每个参数并累加参数数量
    _is_layer = not isinstance(model, torch.nn.modules.container.Sequential)
    if _is_layer:
        layer_params = 0
        layer = model
        for param in layer.parameters():
            layer_params += param.numel()
        # if debug: print('layer_params:', layer_params)
        return layer_params

    total_params = 0
    with torch.no_grad():
        for layer in model:
            # X = layer(X)
            layer_params = 0
            for param in layer.parameters():
                layer_params += param.numel()
            total_params += layer_params
            # if dbg:
            #     print('layer_params:', layer_params, '---', layer.__class__.__name__, 'output shape:\t', X.shape)
    return total_params


if __name__ == '__main__':
    debug = True
    # image_size = [3, 40, 60]
    # X = torch.zeros([2, image_size[0], image_size[1], image_size[2]])

    # _X = torch.ones([3, image_size[0], image_size[1], image_size[2]])
    # X = _X.reshape(3, image_size[0], image_size[1] * image_size[2])
    # shape = X.shape
    #
    # valid_lens = torch.IntTensor([3, 2, 3])
    # valid_lens.dim()
    # # sequence_mask(X, valid_lens, value=-1e6).shape
    #
    # print(sequence_mask(X, valid_lens))
    # X.reshape(-1, shape[-1]).shape
    # masked_softmax(X, valid_lens)

    # qkv = torch.ones((2, 2, 100))

    image_size = [3, 40, 60]
    num_hiddens = image_size[1] * image_size[2]
    batch_size = 2
    # qkv = torch.ones((batch_size, image_size[0], num_hiddens))
    qkv = torch.ones([batch_size, image_size[0], image_size[1], image_size[2]])
    valid_lens = None

    # valid_lens = torch.IntTensor([1] * batch_size)
    # # masked_softmax(qkv, valid_lens)
    # attention = DotProductAttention(dropout=0.5)
    # res = attention(qkv, qkv, qkv, valid_lens)

    # --- multi-head
    num_heads = 1
    # key_size, query_size, value_size, num_hiddens, num_heads, dropout
    attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens, num_hiddens, num_heads, dropout=0.5)
    # attention.eval()
    get_total_parameters(attention)

    res = attention(qkv, qkv, qkv, valid_lens)
    print(qkv.shape, '-->', res.shape)

    if 1:
        old_shape = qkv.shape
        qkv = qkv.reshape([batch_size, image_size[0], -1])
        qkv.shape
        res = attention(qkv, qkv, qkv, valid_lens)
        res.reshape(old_shape).shape
        # get_total_parameters(attention)

    # get_total_parameters(attention, is_layer=False)
    # net = nn.Sequential(attention, attention)
    # get_total_parameters(net)

    batch_size, num_queries = 2, 4
    num_kvpairs, valid_lens = 6, torch.tensor([3, 2])
    X = torch.ones((batch_size, num_queries, num_hiddens))
    Y = torch.ones((batch_size, num_kvpairs, num_hiddens))
    print('shape 0 ------ X[b, q, hiddens], Y[b, kv, hiddens]:', X.shape, Y.shape)
    attention(X, Y, Y, valid_lens).shape

    1